import random
import copy
import cv2
import numpy as np
from PIL import Image, ImageColor, ImageFont, ImageDraw, ImageFilter
import re
import os
import tqdm
import shutil
from tool import filesystem, opencv_tool


def split_char_to_a_and_b(text_color="#000000,#202020"):
    """
    将字符通过字体文件进行分级保存
    """
    char_txt = "/home/swls/work_dir/ocr/synth/TextRecognitionDataGenerator-master_one/TextRecognitionDataGenerator/char_A.txt"
    font_path = "/home/swls/work_dir/ocr/synth/TextRecognitionDataGenerator-master_one/TextRecognitionDataGenerator/fonts/common_ocr/font/FZBIAOYSB.TTF"

    char_a = []
    char_b = []

    image_font = ImageFont.truetype(font=font_path, size=64)

    with open(char_txt, "r", encoding="utf8") as rf:

        for line in rf.readlines():
            char = line.replace("\n", "")

            text_width, text_height = image_font.getsize(char)
            txt_img = Image.new('RGBA', (text_width , text_height), (255, 255, 255, 0))
            txt_draw = ImageDraw.Draw(txt_img)
            colors = [ImageColor.getrgb(c) for c in text_color.split(',')]
            c1, c2 = colors[0], colors[-1]
            fill = (
                random.randint(c1[0], c2[0]),
                random.randint(c1[1], c2[1]),
                random.randint(c1[2], c2[2])
            )
            txt_draw.text((0, 0), char, fill=fill, font=image_font)

            final_image = txt_img.convert('RGBA')
            image_array = np.frombuffer(final_image.tobytes(), dtype=np.uint8)
            image_array = image_array.reshape((final_image.size[1], final_image.size[0], 4))
            cv_image = copy.copy(image_array[:,:,:3])

            if np.mean(cv_image) == 255:
                char_b.append(char)
            else:
                char_a.append(char)
            
        print(char_b)
        print("len(char_b): ", len(char_b))
        print("len(char_a): ", len(char_a))
        
        with open("a.txt", "w", encoding="utf8") as wf:
            for a in char_a:
                wf.write(a + "\n")

        with open("b.txt", "w", encoding="utf8") as wf:
            for b in char_b:
                wf.write(b + "\n")


def group_font_by_image():
    """
    根据手工分类好的字体图像文件, 匹配字体文件并保存目录结构
    """

    image_dir = "/home/swls/work_dir/ocr/synth/TextRecognitionDataGenerator-master_one/TextRecognitionDataGenerator/test_fonts_img"
    font_dir = "/home/swls/work_dir/ocr/synth/TextRecognitionDataGenerator-master_one/TextRecognitionDataGenerator/fonts/common_ocr/cn"
    dest_save_root = "/home/swls/work_dir/ocr/synth/TextRecognitionDataGenerator-master_one/TextRecognitionDataGenerator/fonts/common_ocr"

    all_image_path = filesystem.get_all_filepath(image_dir, [".jpg"])

    for image_path in all_image_path:
        font_name = os.path.basename(image_path).replace(".jpg", "")
        save_dir = dest_save_root + os.sep + os.path.dirname(image_path).replace(image_dir, "") 
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)
        save_path = save_dir + os.sep + font_name
        shutil.copy(font_dir+ os.sep + font_name, save_path)


def move_font_by_img(image_dir, font_dir, dest_save_root, copy=False, save_dir_struct=False):
    """
    通过图像名称匹配字体并移动到指定文件夹
    手工筛选不符合的字体并更新位置分类
    @copy: 是否执行复制 还是移动
    @save_dir_struct: 是否保存目录结构
    """

    if not os.path.exists(dest_save_root):
        os.makedirs(dest_save_root)

    all_image_path = filesystem.get_all_filepath(image_dir, [".jpg"])
    all_font_path = filesystem.get_all_filepath(font_dir, ["ttf","TTF" , "otf","OTF", "ttc", "TTC"])

    for image_path in all_image_path:
        match_ok = False

        for font_path in all_font_path:
            image_name = os.path.basename(image_path).replace(".jpg", "")
            # font_name = os.path.basename(image_path).replace(".jpg", "")
            font_name = os.path.basename(font_path)

            if image_name == font_name:
                match_ok = True
                if save_dir_struct:
                    save_path_dir = dest_save_root + os.path.dirname(font_path).replace(font_dir, "")
                    os.makedirs(save_path_dir, exist_ok=True)
                    save_path = save_path_dir  + os.sep + font_name
                else:
                    save_path = dest_save_root + os.sep + font_name

                if copy:
                    shutil.copy(font_path, save_path)
                else:
                    shutil.move(font_path, save_path)
                break

        if not match_ok:
            print("not match... ", image_path)

def move_image_by_mean():
    """
    根据图像均值来区分字体的等级
    """
    image_dir = "/home/swls/work_dir/ocr/synth/TextRecognitionDataGenerator-master_one/TextRecognitionDataGenerator/test_fonts_img"
    save_a = image_dir + os.sep + "A"
    save_b = image_dir + os.sep + "B"
    save_c = image_dir + os.sep + "C"
    save_d = image_dir + os.sep + "D"

    for n in os.listdir(image_dir):
        image_path = image_dir + os.sep + n

        # if n != "FZXIYSB.TTF.jpg":
        #     continue
        if not os.path.isfile(image_path):
            continue
        
        image = cv2.imread(image_path)
        if image is None:
            print(image_path)

        if np.mean(image[:, image.shape[1] * 500 // 1958:]) == 255:   # a
            shutil.copy(image_path, save_a + os.sep + os.path.basename(image_path) )
        elif np.mean(image[:, image.shape[1] * 1000 // 1958:]) == 255:   # b
            shutil.copy(image_path, save_b + os.sep + os.path.basename(image_path) )
        elif np.mean(image[:, image.shape[1] * 1500 // 1958: ]) == 255:   # c
            shutil.copy(image_path, save_c + os.sep + os.path.basename(image_path) )
        else :   # d
            shutil.copy(image_path, save_d + os.sep + os.path.basename(image_path) )
 

def clear_all_font_images():
    """
    删除图像，保存目录结构
    """
    image_dir = "/home/swls/work_dir/ocr/synth/TextRecognitionDataGenerator-master_one/TextRecognitionDataGenerator/test_fonts_img"

    all_image_path = filesystem.get_all_filepath(image_dir, [".jpg"])

    for image_path in all_image_path:
        os.remove(image_path)


def check_font_and_char_plus():
    """
    检查字体是否支持所有字符
    每张图像x行，每行40个字
    """

    per_char_in_line = 40

    font_dir = "/home/swls/work_dir/ocr/synth/TextRecognitionDataGenerator-master_one/TextRecognitionDataGenerator/fonts/common_ocr/C/print"
    char_path = "/home/swls/work_dir/ocr/synth/TextRecognitionDataGenerator-master_one/TextRecognitionDataGenerator/char_A.txt"
    dest_save_root = "/home/swls/work_dir/ocr/synth/TextRecognitionDataGenerator-master_one/TextRecognitionDataGenerator/test_fonts_img"

    char_list = [c.replace("\n", "") for c in open(char_path, "r", encoding="utf8").readlines()]
    char_length = len(char_list)
    print("char_list: ", char_length)
    
    all_font_path = filesystem.get_all_filepath(font_dir, [".ttf", ".TTF", ".OTF", ".otf"])

    for font_path in tqdm.tqdm(all_font_path, total=len(all_font_path)):

        image_font = ImageFont.truetype(font=font_path, size=64)

        one_image = None
        line_image = None

        for char_idx in range(char_length):
        

            char = char_list[char_idx]
            text_width, text_height = image_font.getsize(char)
            txt_img = Image.new('RGBA', (text_width , text_height), (255, 255, 255, 0))
            txt_draw = ImageDraw.Draw(txt_img)
            
            txt_draw.text((0, 0), char, fill=(0, 0,0), font=image_font)

            final_image = txt_img.convert('RGBA')
            image_array = np.frombuffer(final_image.tobytes(), dtype=np.uint8)
            image_array = image_array.reshape((final_image.size[1], final_image.size[0], 4))
            cv_image = copy.copy(image_array[:,:,:3])

            if char_idx != 0 and char_idx % per_char_in_line == 0:
                if one_image is None:
                    one_image = copy.copy(line_image)
                else:
                    one_image = opencv_tool.vconcatenate_image(one_image,line_image,0 )
                line_image = None
            else:
                if line_image is None:
                    line_image = cv_image
                else:
                    line_image = opencv_tool.hconcatenate_image(line_image,cv_image,0 )
        
        one_image = opencv_tool.vconcatenate_image(one_image,line_image,0 )
        image_save_path = dest_save_root + os.sep + os.path.basename(font_path) + ".jpg"
        cv2.imwrite(image_save_path, one_image)



def check_font_and_char():
    """
    检查字体是否支持所有字符
    每张图像15行，每行30个字
    """

    per_line_in_image = 25
    per_char_in_line = 40

    font_dir = "/home/swls/work_dir/ocr/synth/TextRecognitionDataGenerator-master_one/TextRecognitionDataGenerator/fonts/common_ocr/D/print"
    char_path = "/home/swls/work_dir/ocr/synth/TextRecognitionDataGenerator-master_one/TextRecognitionDataGenerator/char_D.txt"
    dest_save_root = "/home/swls/work_dir/ocr/synth/TextRecognitionDataGenerator-master_one/TextRecognitionDataGenerator/test_fonts_img"

    char_list = [c.replace("\n", "") for c in open(char_path, "r", encoding="utf8").readlines()]
    char_length = len(char_list)
    print("char_list: ", char_length)
    
    all_font_path = filesystem.get_all_filepath(font_dir, [".ttf", ".TTF", ".OTF", ".otf"])

    image_nums = char_length // (per_line_in_image * per_char_in_line) 
    if char_length % (per_line_in_image * per_char_in_line) != 0:
        image_nums += 1

    for font_path in tqdm.tqdm(all_font_path, total=len(all_font_path)):

        image_font = ImageFont.truetype(font=font_path, size=64)


        for image_idx in range(image_nums):
            
            one_image = None
            line_image = None

            for char_idx in range(image_idx*(per_line_in_image * per_char_in_line), (image_idx+1)*(per_line_in_image * per_char_in_line)):
            
                if char_idx >= char_length:
                    continue
                
                char = char_list[char_idx]
                text_width, text_height = image_font.getsize(char)
                txt_img = Image.new('RGBA', (text_width , text_height), (255, 255, 255, 0))
                txt_draw = ImageDraw.Draw(txt_img)
                

                txt_draw.text((0, 0), char, fill=(0, 0,0), font=image_font)

                final_image = txt_img.convert('RGBA')
                image_array = np.frombuffer(final_image.tobytes(), dtype=np.uint8)
                image_array = image_array.reshape((final_image.size[1], final_image.size[0], 4))
                cv_image = copy.copy(image_array[:,:,:3])

                if char_idx != 0 and char_idx % per_char_in_line == 0:
                    if one_image is None:
                        one_image = copy.copy(line_image)
                    else:
                        one_image = opencv_tool.vconcatenate_image(one_image,line_image,0 )
                    line_image = None

                if line_image is None:
                    line_image = cv_image
                else:
                    line_image = opencv_tool.hconcatenate_image(line_image,cv_image,0 )
        
            one_image = opencv_tool.vconcatenate_image(one_image,line_image,0 )
            image_save_path = dest_save_root + os.sep + os.path.basename(font_path) + ".{}.jpg".format(char_idx)
            cv2.imwrite(image_save_path, one_image)

def rename_batch(file_dir):
    all_files = filesystem.get_all_filepath(file_dir, [".ttf", ".TTF", ".OTF", ".otf", "TTC", "ttc"])
    for path in all_files:
        new_path = path + ".jpg"
        shutil.move(path, new_path)

if __name__ == "__main__":
    # split_char_to_a_and_b()
    
    # group_font_by_image()

    image_dir = "/home/swls/work_dir/ocr/synth/TextRecognitionDataGenerator-master_one/test_fonts_img/1"
    font_dir = "/home/swls/work_dir/ocr/synth/TextRecognitionDataGenerator-master_one/fonts/meng_ch/D/print"
    dest_save_root = "/home/swls/work_dir/ocr/synth/TextRecognitionDataGenerator-master_one/fonts/meng_ch/D/1"
    move_font_by_img(image_dir, font_dir, dest_save_root, copy=True, save_dir_struct=True)

    # file_dir = "/home/swls/work_dir/ocr/synth/TextRecognitionDataGenerator-master_one/TextRecognitionDataGenerator/fonts/driving_license/tmp"
    # rename_batch(file_dir)

    # move_image_by_mean()

    # clear_all_font_images()

    # check_font_and_char()
    # check_font_and_char_plus()

    